/*
 * Copyright 2022 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package androidx.compose.foundation.gestures.snapping

import androidx.compose.animation.core.AnimationScope
import androidx.compose.animation.core.AnimationSpec
import androidx.compose.animation.core.AnimationState
import androidx.compose.animation.core.AnimationVector
import androidx.compose.animation.core.AnimationVector1D
import androidx.compose.animation.core.DecayAnimationSpec
import androidx.compose.animation.core.Spring
import androidx.compose.animation.core.animateDecay
import androidx.compose.animation.core.animateTo
import androidx.compose.animation.core.calculateTargetValue
import androidx.compose.animation.core.copy
import androidx.compose.animation.core.spring
import androidx.compose.animation.rememberSplineBasedDecay
import androidx.compose.foundation.ExperimentalFoundationApi
import androidx.compose.foundation.gestures.DefaultScrollMotionDurationScale
import androidx.compose.foundation.gestures.FlingBehavior
import androidx.compose.foundation.gestures.ScrollScope
import androidx.compose.foundation.gestures.TargetedFlingBehavior
import androidx.compose.foundation.internal.checkPrecondition
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.unit.dp
import kotlin.coroutines.cancellation.CancellationException
import kotlin.math.abs
import kotlin.math.absoluteValue
import kotlin.math.sign
import kotlinx.coroutines.withContext

/**
 * A [TargetedFlingBehavior] that performs snapping to a given position in a layout.
 *
 * You can use [SnapLayoutInfoProvider.calculateApproachOffset] to indicate that snapping should
 * happen after this offset. If the velocity generated by the fling is high enough to get there,
 * we'll use [decayAnimationSpec] to get to that offset and then we'll snap to the next bound
 * calculated by [SnapLayoutInfoProvider.calculateSnapOffset] using [snapAnimationSpec].
 *
 * If the velocity is not high enough, we'll use [snapAnimationSpec] to approach and the same spec
 * to snap into place.
 *
 * Please refer to the sample to learn how to use this API.
 *
 * @sample androidx.compose.foundation.samples.SnapFlingBehaviorSimpleSample
 * @sample androidx.compose.foundation.samples.SnapFlingBehaviorCustomizedSample
 * @param snapLayoutInfoProvider The information about the layout being snapped.
 * @param decayAnimationSpec The animation spec used to approach the target offset when the fling
 *   velocity is large enough. Large enough means large enough to naturally decay.
 * @param snapAnimationSpec The animation spec used to finally snap to the correct bound.
 */
fun snapFlingBehavior(
    snapLayoutInfoProvider: SnapLayoutInfoProvider,
    decayAnimationSpec: DecayAnimationSpec<Float>,
    snapAnimationSpec: AnimationSpec<Float>,
): TargetedFlingBehavior {
    return SnapFlingBehavior(snapLayoutInfoProvider, decayAnimationSpec, snapAnimationSpec)
}

internal class SnapFlingBehavior(
    private val snapLayoutInfoProvider: SnapLayoutInfoProvider,
    private val decayAnimationSpec: DecayAnimationSpec<Float>,
    private val snapAnimationSpec: AnimationSpec<Float>,
) : TargetedFlingBehavior {

    internal var motionScaleDuration = DefaultScrollMotionDurationScale

    /**
     * Perform a snapping fling animation with given velocity and suspend until fling has finished.
     * This will behave the same way as [performFling] except it will report on each
     * remainingOffsetUpdate using the [onRemainingDistanceUpdated] lambda.
     *
     * @param initialVelocity velocity available for fling in the orientation specified in
     *   [androidx.compose.foundation.gestures.scrollable] that invoked this method.
     * @param onRemainingDistanceUpdated a lambda that will be called anytime the distance to the
     *   settling offset is updated. The settling offset is the final offset where this fling will
     *   stop and may change depending on the snapping animation progression.
     * @return remaining velocity after fling operation has ended
     */
    override suspend fun ScrollScope.performFling(
        initialVelocity: Float,
        onRemainingDistanceUpdated: (Float) -> Unit,
    ): Float {
        val (remainingOffset, remainingState) = fling(initialVelocity, onRemainingDistanceUpdated)

        debugLog { "Post Settling Offset=$remainingOffset" }
        // No remaining offset means we've used everything, no need to propagate velocity. Otherwise
        // we couldn't use everything (probably because we have hit the min/max bounds of the
        // containing layout) we should propagate the offset.
        return if (remainingOffset == 0f) NoVelocity else remainingState.velocity
    }

    private suspend fun ScrollScope.fling(
        initialVelocity: Float,
        onRemainingScrollOffsetUpdate: (Float) -> Unit,
    ): AnimationResult<Float, AnimationVector1D> {
        val result =
            withContext(motionScaleDuration) {
                val decayOffset =
                    decayAnimationSpec.calculateTargetValue(
                        initialVelocity = initialVelocity,
                        initialValue = 0.0f,
                    )

                val initialOffset =
                    snapLayoutInfoProvider.calculateApproachOffset(initialVelocity, decayOffset)

                checkPrecondition(!initialOffset.isNaN()) {
                    "calculateApproachOffset returned NaN. Please use a valid value."
                }

                // ensure offset sign and value are correct
                var remainingScrollOffset = abs(initialOffset) * sign(initialVelocity)

                onRemainingScrollOffsetUpdate(remainingScrollOffset) // First Scroll Offset

                val animationState =
                    tryApproach(remainingScrollOffset, initialVelocity) { delta ->
                        remainingScrollOffset -= delta
                        onRemainingScrollOffsetUpdate(remainingScrollOffset)
                    }

                val finalSnapOffset =
                    snapLayoutInfoProvider.calculateSnapOffset(animationState.velocity)

                checkPrecondition(!finalSnapOffset.isNaN()) {
                    "calculateSnapOffset returned NaN. Please use a valid value."
                }

                remainingScrollOffset = finalSnapOffset

                debugLog { "Settling Final Bound=$remainingScrollOffset" }

                animateWithTarget(
                    remainingScrollOffset,
                    remainingScrollOffset,
                    animationState.copy(value = 0f), // re-use the velocity and timestamp from state
                    snapAnimationSpec,
                ) { delta ->
                    remainingScrollOffset -= delta
                    onRemainingScrollOffsetUpdate(remainingScrollOffset)
                }
            }

        onRemainingScrollOffsetUpdate(0f) // Animation finished or was cancelled
        return result
    }

    private suspend fun ScrollScope.tryApproach(
        offset: Float,
        velocity: Float,
        updateRemainingScrollOffset: (Float) -> Unit,
    ): AnimationState<Float, AnimationVector1D> {
        // If we don't have velocity or approach offset, we shouldn't run the approach animation
        return if (offset.absoluteValue == 0.0f || velocity.absoluteValue == 0.0f) {
            AnimationState(offset, velocity)
        } else {
            runApproach(offset, velocity, updateRemainingScrollOffset).currentAnimationState
        }
    }

    private suspend fun ScrollScope.runApproach(
        initialTargetOffset: Float,
        initialVelocity: Float,
        onAnimationStep: (delta: Float) -> Unit,
    ): AnimationResult<Float, AnimationVector1D> {

        val animation =
            if (isDecayApproachPossible(offset = initialTargetOffset, velocity = initialVelocity)) {
                debugLog { "High Velocity Approach" }
                DecayApproachAnimation(decayAnimationSpec)
            } else {
                debugLog { "Low Velocity Approach" }
                TargetApproachAnimation(snapAnimationSpec)
            }

        return approach(initialTargetOffset, initialVelocity, animation, onAnimationStep)
    }

    /** If we can approach the target and still have velocity left */
    private fun isDecayApproachPossible(offset: Float, velocity: Float): Boolean {
        val decayOffset = decayAnimationSpec.calculateTargetValue(NoDistance, velocity)
        debugLog {
            "Evaluating decay possibility with " +
                "decayOffset=$decayOffset and proposed approach=$offset"
        }
        return decayOffset.absoluteValue >= offset.absoluteValue
    }

    @Suppress("Deprecation")
    override fun equals(other: Any?): Boolean {
        return if (other is SnapFlingBehavior) {
            other.snapAnimationSpec == this.snapAnimationSpec &&
                other.decayAnimationSpec == this.decayAnimationSpec &&
                other.snapLayoutInfoProvider == this.snapLayoutInfoProvider
        } else {
            false
        }
    }

    override fun hashCode(): Int =
        0.let { 31 * it + snapAnimationSpec.hashCode() }
            .let { 31 * it + decayAnimationSpec.hashCode() }
            .let { 31 * it + snapLayoutInfoProvider.hashCode() }
}

/**
 * Creates and remember a [FlingBehavior] that performs snapping.
 *
 * @param snapLayoutInfoProvider The information about the layout that will do snapping
 */
@Composable
fun rememberSnapFlingBehavior(
    snapLayoutInfoProvider: SnapLayoutInfoProvider
): TargetedFlingBehavior {
    val density = LocalDensity.current
    val highVelocityApproachSpec: DecayAnimationSpec<Float> = rememberSplineBasedDecay()
    return remember(snapLayoutInfoProvider, highVelocityApproachSpec, density) {
        snapFlingBehavior(
            snapLayoutInfoProvider = snapLayoutInfoProvider,
            decayAnimationSpec = highVelocityApproachSpec,
            snapAnimationSpec = spring(stiffness = Spring.StiffnessMediumLow),
        )
    }
}

/**
 * To ensure we do not overshoot, the approach animation is divided into 2 parts.
 *
 * In the initial animation we animate up until targetOffset. At this point we will have fulfilled
 * the requirement of [SnapLayoutInfoProvider.calculateApproachOffset] and we should snap to the
 * next [SnapLayoutInfoProvider.calculateSnapOffset].
 *
 * The second part of the approach is a UX improvement. If the target offset is too far (in here, we
 * define too far as over half a step offset away) we continue the approach animation a bit further
 * and leave the remainder to be snapped.
 */
private suspend fun ScrollScope.approach(
    initialTargetOffset: Float,
    initialVelocity: Float,
    animation: ApproachAnimation<Float, AnimationVector1D>,
    onAnimationStep: (delta: Float) -> Unit,
): AnimationResult<Float, AnimationVector1D> {

    return animation.approachAnimation(this, initialTargetOffset, initialVelocity, onAnimationStep)
}

private class AnimationResult<T, V : AnimationVector>(
    val remainingOffset: T,
    val currentAnimationState: AnimationState<T, V>,
) {
    operator fun component1(): T = remainingOffset

    operator fun component2(): AnimationState<T, V> = currentAnimationState
}

private operator fun <T : Comparable<T>> ClosedFloatingPointRange<T>.component1(): T = this.start

private operator fun <T : Comparable<T>> ClosedFloatingPointRange<T>.component2(): T =
    this.endInclusive

/**
 * Run a [DecayAnimationSpec] animation up to before [targetOffset] using [animationState]
 *
 * @param targetOffset The destination of this animation. Since this is a decay animation, we can
 *   use this value to prevent the animation to run until the end.
 * @param animationState The previous [AnimationState] for continuation purposes.
 * @param decayAnimationSpec The [DecayAnimationSpec] that will drive this animation
 * @param onAnimationStep Called for each new scroll delta emitted by the animation cycle.
 */
@OptIn(ExperimentalFoundationApi::class)
private suspend fun ScrollScope.animateDecay(
    targetOffset: Float,
    animationState: AnimationState<Float, AnimationVector1D>,
    decayAnimationSpec: DecayAnimationSpec<Float>,
    onAnimationStep: (delta: Float) -> Unit,
): AnimationResult<Float, AnimationVector1D> {
    var previousValue = 0f

    fun AnimationScope<Float, AnimationVector1D>.consumeDelta(delta: Float) {
        var consumed = 0.0f
        try {
            consumed = scrollBy(delta)
        } catch (_: CancellationException) {
            cancelAnimation()
        }

        onAnimationStep(consumed)
        if (abs(delta - consumed) > 0.5f) cancelAnimation()
    }

    animationState.animateDecay(
        decayAnimationSpec,
        sequentialAnimation = animationState.velocity != 0f,
    ) {
        if (abs(value) >= abs(targetOffset)) {
            val finalValue = value.coerceToTarget(targetOffset)
            val finalDelta = finalValue - previousValue
            consumeDelta(finalDelta)
            cancelAnimation()
            previousValue = finalValue
        } else {
            val delta = value - previousValue
            consumeDelta(delta)
            previousValue = value
        }
    }

    debugLog { "Decay Animation: Proposed Offset=$targetOffset Achieved Offset=$previousValue" }
    return AnimationResult(targetOffset - previousValue, animationState)
}

/**
 * Runs a [AnimationSpec] to snap the list into [targetOffset]. Uses [cancelOffset] to stop this
 * animation before it reaches the target.
 *
 * @param targetOffset The final target of this animation
 * @param cancelOffset If we'd like to finish the animation earlier we use this value
 * @param animationState The current animation state for continuation purposes
 * @param animationSpec The [AnimationSpec] that will drive this animation
 * @param onAnimationStep Called for each new scroll delta emitted by the animation cycle.
 */
@OptIn(ExperimentalFoundationApi::class)
private suspend fun ScrollScope.animateWithTarget(
    targetOffset: Float,
    cancelOffset: Float,
    animationState: AnimationState<Float, AnimationVector1D>,
    animationSpec: AnimationSpec<Float>,
    onAnimationStep: (delta: Float) -> Unit,
): AnimationResult<Float, AnimationVector1D> {
    var consumedUpToNow = 0f
    val initialVelocity = animationState.velocity
    animationState.animateTo(
        targetOffset,
        animationSpec = animationSpec,
        sequentialAnimation = (animationState.velocity != 0f),
    ) {
        val realValue = value.coerceToTarget(cancelOffset)
        val delta = realValue - consumedUpToNow
        var consumed = 0.0f
        try {
            consumed = scrollBy(delta)
        } catch (_: CancellationException) {
            cancelAnimation()
        }
        onAnimationStep(consumed)
        // stop when unconsumed or when we reach the desired value
        if (abs(delta - consumed) > 0.5f || realValue != value) {
            cancelAnimation()
        }
        consumedUpToNow += consumed
    }

    debugLog { "Snap Animation: Proposed Offset=$targetOffset Achieved Offset=$consumedUpToNow" }

    // Always course correct velocity so they don't become too large.
    val finalVelocity = animationState.velocity.coerceToTarget(initialVelocity)
    return AnimationResult(
        targetOffset - consumedUpToNow,
        animationState.copy(velocity = finalVelocity),
    )
}

private fun Float.coerceToTarget(target: Float): Float {
    if (target == 0f) return 0f
    return if (target > 0) coerceAtMost(target) else coerceAtLeast(target)
}

/** The animations used to approach offset and approach half a step offset. */
private interface ApproachAnimation<T, V : AnimationVector> {
    suspend fun approachAnimation(
        scope: ScrollScope,
        offset: T,
        velocity: T,
        onAnimationStep: (delta: T) -> Unit,
    ): AnimationResult<T, V>
}

private class TargetApproachAnimation(private val animationSpec: AnimationSpec<Float>) :
    ApproachAnimation<Float, AnimationVector1D> {
    override suspend fun approachAnimation(
        scope: ScrollScope,
        offset: Float,
        velocity: Float,
        onAnimationStep: (delta: Float) -> Unit,
    ): AnimationResult<Float, AnimationVector1D> {
        val animationState = AnimationState(initialValue = 0f, initialVelocity = velocity)
        val targetOffset = offset.absoluteValue * sign(velocity)
        return with(scope) {
            animateWithTarget(
                targetOffset = targetOffset,
                cancelOffset = offset,
                animationState = animationState,
                animationSpec = animationSpec,
                onAnimationStep = onAnimationStep,
            )
        }
    }
}

private class DecayApproachAnimation(private val decayAnimationSpec: DecayAnimationSpec<Float>) :
    ApproachAnimation<Float, AnimationVector1D> {
    override suspend fun approachAnimation(
        scope: ScrollScope,
        offset: Float,
        velocity: Float,
        onAnimationStep: (delta: Float) -> Unit,
    ): AnimationResult<Float, AnimationVector1D> {
        val animationState = AnimationState(initialValue = 0f, initialVelocity = velocity)
        return with(scope) {
            animateDecay(offset, animationState, decayAnimationSpec, onAnimationStep)
        }
    }
}

internal val MinFlingVelocityDp = 400.dp
internal const val NoDistance = 0f
internal const val NoVelocity = 0f

internal fun calculateFinalOffset(
    snappingOffset: FinalSnappingItem,
    lowerBound: Float,
    upperBound: Float,
): Float {

    fun Float.isValidDistance(): Boolean {
        return this != Float.POSITIVE_INFINITY && this != Float.NEGATIVE_INFINITY
    }

    debugLog { "Proposed Bounds: Lower=$lowerBound Upper=$upperBound" }

    val finalDistance =
        when (snappingOffset) {
            FinalSnappingItem.ClosestItem -> {
                if (abs(upperBound) <= abs(lowerBound)) {
                    upperBound
                } else {
                    lowerBound
                }
            }
            FinalSnappingItem.NextItem -> upperBound
            FinalSnappingItem.PreviousItem -> lowerBound
            else -> NoDistance
        }

    return if (finalDistance.isValidDistance()) {
        finalDistance
    } else {
        NoDistance
    }
}

private const val DEBUG = false

private inline fun debugLog(generateMsg: () -> String) {
    if (DEBUG) {
        println("SnapFlingBehavior: ${generateMsg()}")
    }
}
